# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""

import torch
import torchvision
from PIL import Image
from torchvision import transforms

from utils import pytorch_utils as utils
from vgg16_net import Vgg16

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    # net: torch.nn.Module = torch.hub.load("pytorch/vision:v0.10.0", 'vgg11', pretrained=True)
    device = utils.get_device()
    state_dict = torch.load("out/vgg16_9_20230918101012.pth", map_location=device)

    net: torch.nn.Module = Vgg16()
    # net: torch.nn.Module = torchvision.models.vgg16()
    # net.classifier[6] = torch.nn.Linear(4096, 10)
    net.load_state_dict(state_dict)
    net.to(device)
    net.eval()

    image = Image.open("../../assets/cat.jpg")

    transform = transforms.Compose(
        [
            transforms.Resize(33),
            transforms.CenterCrop(32),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]
    )

    image = transform(image)
    image = image.unsqueeze(0)
    image = image.to(device)

    with torch.no_grad():
        out = net(image)

    classes = ('plane', 'car', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck')

    probabilities = torch.nn.functional.softmax(out[0], dim=0)
    print("probabilities:", probabilities)
    top5_prob, top5_catid = torch.topk(probabilities, 5)
    for i in range(top5_prob.size(0)):
        print("===================")
        catId = top5_catid[i].item()
        print("top5_prob:", top5_prob[i])
        print("top5_catid:", catId)
        print("top5_catss:", classes[catId])
